"""Cookie加密存储模块。"""

import base64
import os
from typing import Optional
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import logging

logger = logging.getLogger(__name__)


class CookieEncryption:
    """Cookie加密器。"""
    
    def __init__(self, password: str = None):
        """初始化加密器。
        
        Args:
            password: 主密码，如果为None则从环境变量获取
        """
        if password is None:
            password = os.getenv('CRAWLER_ENCRYPTION_KEY', 'default_key_change_in_production')
            
        # 使用PBKDF2生成密钥
        salt = b'crawler_cookie_salt'  # 在生产环境中应该使用随机盐
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        self.cipher_suite = Fernet(key)
        
    def encrypt_cookie(self, cookie: str) -> str:
        """加密Cookie。
        
        Args:
            cookie: 原始Cookie字符串
            
        Returns:
            加密后的Cookie字符串
        """
        try:
            encrypted_data = self.cipher_suite.encrypt(cookie.encode())
            return base64.urlsafe_b64encode(encrypted_data).decode()
        except Exception as e:
            logger.error(f"Cookie加密失败: {e}")
            raise
            
    def decrypt_cookie(self, encrypted_cookie: str) -> Optional[str]:
        """解密Cookie。
        
        Args:
            encrypted_cookie: 加密的Cookie字符串
            
        Returns:
            解密后的Cookie字符串，失败返回None
        """
        try:
            encrypted_data = base64.urlsafe_b64decode(encrypted_cookie.encode())
            decrypted_data = self.cipher_suite.decrypt(encrypted_data)
            return decrypted_data.decode()
        except Exception as e:
            logger.error(f"Cookie解密失败: {e}")
            return None
            
    def is_encrypted(self, cookie: str) -> bool:
        """检查Cookie是否已加密。
        
        Args:
            cookie: Cookie字符串
            
        Returns:
            是否已加密
        """
        try:
            # 尝试base64解码，如果成功且能解密则认为是加密的
            base64.urlsafe_b64decode(cookie.encode())
            return True
        except Exception:
            return False


# 全局加密器实例
cookie_encryptor = CookieEncryption()